import comfy
import comfy.sd
from comfy import *

# 设计一个父类
class AC_FUN:
     RETURN_TYPES = None
     RETURN_NAMES = None
     FUNCTION = None,
     CATEGORY  = "♋AC_FUNV3.1"
# 模型开关
class Switch_Model(AC_FUN):
    list_1 = ["Model_1", "Model_2"]
    @classmethod
    def INPUT_TYPES(self):
         return {
            "required": {
                 "Model_1":("MODEL",),
                 "CLIP_1" : ("CLIP",),
                 "VAE_1":("VAE",),
                 "Model_2":("MODEL",),
                 "CLIP_2" : ("CLIP",),
                 "VAE_2":("VAE",),
                 "Select" :(self.list_1,)
            }
         }
    RETURN_TYPES = ("MODEL", "CLIP", "VAE")
    RETURN_NAMES = ("MODEL", "CLIP", "VAE")
    FUNCTION = "switch_model"

    def switch_model(self,Model_1,CLIP_1,VAE_1,Model_2,CLIP_2,VAE_2,Select):
         if Select == "Model_1":
              return (Model_1,CLIP_1,VAE_1)
         if Select == "Model_2":
              return (Model_2,CLIP_2,VAE_2)
         

# Lora开关
class Lora_Switch(AC_FUN):
     list_2 = ["lora_1", "lora_2"]
     @classmethod
     def INPUT_TYPES(self):
          return {"required":{
               "Lora_1":("MODEL",),
               "Lcp_1":("CLIP",),
               "Lora_2":("MODEL",),
               "Lcp_2":("CLIP",),
               "Select" :(self.list_2,)
          }}
     
     RETURN_TYPES = ("MODEL", "CLIP")
     RETURN_NAMES = ("LORA", "CLIP")
     FUNCTION = "switch_lora"

     def switch_lora(self,Lora_1,Lcp_1,Lora_2,Lcp_2,Select):
         if Select == "lora_1":
              return (Lora_1,Lcp_1,)
         if Select == "lora_2":
              return (Lora_2,Lcp_2,)

# Conditioning开关
class Condition_Switch(AC_FUN):
     list_3 = ["Conditioning_1", "Conditioning_2"]
     @classmethod
     def INPUT_TYPES(self):
          return {"required":{
               "Conditioning_1":("CONDITIONING",),
               "Conditioning_2":("CONDITIONING",),
               "Select" :(self.list_3,)
          }}
     
     RETURN_TYPES = ("CONDITIONING", )
     RETURN_NAMES = ("Prompt", )
     FUNCTION = "switch_con"

     def switch_con(self,Conditioning_1,Conditioning_2,Select):
         if Select == "Conditioning_1":
              return (Conditioning_1,)
         if Select == "Conditioning_2":
              return (Conditioning_2,)
         
# Conditioning开关
class Condition_P_N_Switch(AC_FUN):
     list_3 = ["Conditioning_1", "Conditioning_2"]
     @classmethod
     def INPUT_TYPES(self):
          return {"required":{
               "positive_1":("CONDITIONING",),
               "negative_1":("CONDITIONING",),
               "positive_2":("CONDITIONING",),
               "negative_2":("CONDITIONING",),
               "Select" :(self.list_3,)
          }}
     
     RETURN_TYPES = ("CONDITIONING","CONDITIONING")
     RETURN_NAMES = ("positive", "negative" )
     FUNCTION = "switch_con_p"

     def switch_con_p(self,positive_1,negative_1,positive_2,negative_2,Select):
         if Select == "Conditioning_1":
              return (positive_1,negative_1)
         if Select == "Conditioning_2":
              return (positive_2,negative_2)

# 词汇开关
class Switch_prompt(AC_FUN):
    Switch_prompt = ["prompt_1", "prompt_2"]
    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
            "Switch": (self.Switch_prompt,),
            "Prompt_1":("STRING",{"forceInput":True}),
            "Prompt_2":("STRING",{"forceInput":True}),
            }
            }
    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("Prompt",)
    FUNCTION = "switch_prompt"

    def switch_prompt(self, Switch,Prompt_1, Prompt_2):
        if Switch == "prompt_1":
            return (Prompt_1,)
        if Switch == "prompt_2":
            return (Prompt_2,)

# 采样器开关
class Switch_ksampler(AC_FUN):
    Switch_ksampler_list = ["ksampler_1", "ksampler_2"]
    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
            "ksampler_1":(comfy.samplers.KSampler.SAMPLERS,{"default":"lcm"} ),
            "ksampler_2":(comfy.samplers.KSampler.SAMPLERS,{"default":"euler_ancestral"}),
            "Switch": (self.Switch_ksampler_list,)
            }
            }
    RETURN_TYPES = ("SAMPLERS",)
    RETURN_NAMES = ("Sampler_Name",)
    FUNCTION = "switch_sampler"

    def switch_sampler(self, Switch,ksampler_1,ksampler_2):
        if Switch == "ksampler_1":
            return (ksampler_1,)
        if Switch == "ksampler_2":
            return (ksampler_2,)
# VAE开关
class Switch_Vae(AC_FUN):
    
    Switch_Vae_list = ["vae_1", "vae_2"]
    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
            "vae_1":("VAE",),
            "vae_2":("VAE",),
            "Switch": (self.Switch_Vae_list,),
            }
            }
    RETURN_TYPES = ("VAE",)
    RETURN_NAMES = ("VAE",)
    FUNCTION = "switch_vae"

    def switch_vae(self, vae_1,vae_2,Switch):
        if Switch == 'vae_1':
            return (vae_1,)
        if Switch == 'vae_2':
            return (vae_2,)

# Latent开关
class Switch_Latent(AC_FUN):
    Switch_Latent_list = ["Latent_1", "Latent_2"]
    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
            "Latent_1":("LATENT",),
            "Latent_2":("LATENT",),
            "Switch": (self.Switch_Latent_list,),
            }
            }
    RETURN_TYPES = ("LATENT",)
    RETURN_NAMES = ("LATENT",)
    FUNCTION = "switch_latent"

    def switch_latent(self, Latent_1,Latent_2,Switch):
        if Switch == 'Latent_1':
            return (Latent_1,)
        if Switch == 'Latent_2':
            return (Latent_2,)
# dispatch分派
class Switch_Dispatch(AC_FUN):

    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
            "bool": ("BOOLEAN", {"default":False}),
            }
            }
    RETURN_TYPES = ("Any",)
    OUTPUT_LIST = True
    RETURN_NAMES = ("Any",)
    FUNCTION = "switch_dispatch"

    def switch_dispatch(self, bool,):
        if bool == True:
            return (True,)
        if bool == False:
            return (False,)
# 图片分派
class Switch_Picture(AC_FUN):
    Picture_list = ['Picture_1', 'Picture_2']
    @classmethod
    def INPUT_TYPES(self):
        return {
            "required": {
            "Picture_1":("IMAGE",{"enable":False}),
            "Picture_2":("IMAGE",),
            "Switch": (self.Picture_list,),
            }
            }
    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("Picture",)
    FUNCTION = "switch_picture"

    def switch_picture(self, Picture_1, Picture_2, Switch):
        if Switch == 'Picture_1':
            return (Picture_1,)
        if Switch == 'Picture_2':
            return (Picture_2,)



NODE_CLASS_MAPPINGS = {
     "AC_Model_Switch": Switch_Model,
     "AC_Lora_Select": Lora_Switch,
     "AC_COND_Select": Condition_Switch,
     "AC_Condition_P_N_Switch":Condition_P_N_Switch,
     "AC_Prompt_Switch": Switch_prompt,
     "AC_KSampler_Switch": Switch_ksampler,
     "AC_Vae_Switch": Switch_Vae,
     "AC_Latent_Switch": Switch_Latent,
     "AC_Dispatch_Switch": Switch_Dispatch,
     "AC_Picture_Switch": Switch_Picture,
}